# -*- coding: utf-8 -*-
# @File    : cookies_manager.py
# 多Cookie管理器，实现Cookie的加载、轮询分配和任务分片功能

import os
import glob
import random
import time
from typing import List, Dict, Optional, Tuple
import json

class CookiesManager:
    """
    Cookie管理器，用于管理多个Cookie，支持轮询分配和任务分片
    """
    
    def __init__(self, cookies_dir: str = 'cookies'):
        """
        初始化Cookie管理器
        
        Args:
            cookies_dir: Cookie文件所在目录，默认为'cookies'
        """
        self.cookies_dir = cookies_dir
        self.cookies = []  # 存储所有加载的cookie
        self.cookies_status = {}  # 存储cookie状态，如是否可用、剩余请求次数等
        self.load_all_cookies()
    
    def load_all_cookies(self) -> None:
        """
        加载cookies目录下的所有cookie文件
        """
        # 确保cookies目录存在
        if not os.path.exists(self.cookies_dir):
            raise FileNotFoundError(f"Cookie目录 {self.cookies_dir} 不存在")
        
        # 获取cookies目录下的所有.txt文件
        cookie_files = glob.glob(os.path.join(self.cookies_dir, "*.txt"))
        
        if not cookie_files:
            raise FileNotFoundError(f"在 {self.cookies_dir} 目录下未找到任何cookie文件")
        
        # 加载每个cookie文件
        for file_path in cookie_files:
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    cookie_content = f.read().strip()
                    
                # 保存cookie内容和相关状态信息
                cookie_id = os.path.basename(file_path).split('.')[0]  # 文件名作为cookie ID
                self.cookies.append({
                    'id': cookie_id,
                    'content': cookie_content,
                    'file_path': file_path
                })
                
                # 初始化cookie状态
                self.cookies_status[cookie_id] = {
                    'available': True,  # 是否可用
                    'last_used': 0,     # 上次使用时间
                    'request_count': 0  # 请求次数
                }
                
                print(f"已加载Cookie: {cookie_id}")
            except Exception as e:
                print(f"加载Cookie文件 {file_path} 失败: {str(e)}")
        
        print(f"总共加载了 {len(self.cookies)} 个Cookie")
    
    def get_cookie(self, cookie_id: str = None) -> Optional[str]:
        """
        获取指定ID的cookie内容
        
        Args:
            cookie_id: cookie的ID，如果为None则随机选择一个
            
        Returns:
            cookie内容，如果找不到则返回None
        """
        if not self.cookies:
            return None
        
        if cookie_id is None:
            # 随机选择一个可用的cookie
            available_cookies = [c for c in self.cookies 
                                if self.cookies_status[c['id']]['available']]
            if not available_cookies:
                return None
            cookie = random.choice(available_cookies)
        else:
            # 查找指定ID的cookie
            cookie = next((c for c in self.cookies if c['id'] == cookie_id), None)
            if cookie is None or not self.cookies_status[cookie['id']]['available']:
                return None
        
        # 更新cookie使用状态
        self.cookies_status[cookie['id']]['last_used'] = time.time()
        self.cookies_status[cookie['id']]['request_count'] += 1
        
        return cookie['content']
    
    def get_all_cookies(self) -> List[Dict]:
        """
        获取所有cookie的信息
        
        Returns:
            包含所有cookie信息的列表
        """
        return self.cookies
    
    def get_cookie_count(self) -> int:
        """
        获取可用cookie的数量
        
        Returns:
            可用cookie数量
        """
        return len([c for c in self.cookies if self.cookies_status[c['id']]['available']])
    
    def split_tasks(self, task_ids: List, output_dir: str = 'Weibo/爬取进度/multi_cookies') -> Dict[str, List]:
        """
        根据cookie数量将任务ID列表分片
        
        Args:
            task_ids: 任务ID列表
            output_dir: 分片任务保存目录
            
        Returns:
            每个cookie对应的任务ID列表的字典
        """
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        cookie_count = self.get_cookie_count()
        if cookie_count == 0:
            raise ValueError("没有可用的Cookie")
        
        # 洗牌task_ids以随机分配
        random.shuffle(task_ids)
        
        # 计算每个cookie应分配的任务数
        avg_tasks_per_cookie = len(task_ids) // cookie_count
        remainder = len(task_ids) % cookie_count
        
        # 分配任务
        task_splits = {}
        start_idx = 0
        
        for i, cookie in enumerate(self.cookies):
            cookie_id = cookie['id']
            if not self.cookies_status[cookie_id]['available']:
                continue
                
            # 计算当前cookie应分配的任务数
            tasks_for_this_cookie = avg_tasks_per_cookie + (1 if i < remainder else 0)
            end_idx = start_idx + tasks_for_this_cookie
            
            # 分配任务
            tasks = task_ids[start_idx:end_idx]
            task_splits[cookie_id] = tasks
            
            # 保存分片任务到文件
            task_file = os.path.join(output_dir, f"tasks_{cookie_id}.json")
            with open(task_file, 'w', encoding='utf-8') as f:
                json.dump(tasks, f)
            
            start_idx = end_idx
            
            print(f"Cookie {cookie_id} 分配了 {len(tasks)} 个任务")
        
        return task_splits
    
    def mark_cookie_unavailable(self, cookie_id: str) -> None:
        """
        标记指定cookie为不可用状态
        
        Args:
            cookie_id: cookie的ID
        """
        if cookie_id in self.cookies_status:
            self.cookies_status[cookie_id]['available'] = False
            print(f"Cookie {cookie_id} 已标记为不可用")
    
    def get_cookie_with_header(self, cookie_id: str = None) -> Tuple[str, Dict]:
        """
        获取cookie及对应的header
        
        Args:
            cookie_id: cookie的ID，如果为None则随机选择一个
            
        Returns:
            (cookie内容, 包含该cookie的header字典)
        """
        cookie_content = self.get_cookie(cookie_id)
        if not cookie_content:
            return None, None
        
        # 构建包含cookie的header
        header = {
            'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/104.0.0.0 Safari/537.36',
            'accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9',
            'accept-encoding': 'gzip, deflate, br',
            'cookie': cookie_content,
            'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8',
            'connection': 'keep-alive'
        }
        
        return cookie_content, header 